from flask import Blueprint, request, jsonify, send_file
from src.models.database import db, PredictionTask, PredictionResult
import pandas as pd
import io
import json
from datetime import datetime

export_bp = Blueprint('export', __name__)

@export_bp.route('/prediction/<task_id>/excel', methods=['GET'])
def export_prediction_to_excel(task_id):
    """导出预测结果到Excel"""
    try:
        # 获取任务信息
        task = PredictionTask.query.filter_by(task_id=task_id).first()
        if not task:
            return jsonify({'code': 404, 'message': '任务不存在'}), 404
        
        # 获取预测结果
        results = PredictionResult.query.filter_by(task_id=task_id).all()
        if not results:
            return jsonify({'code': 404, 'message': '没有找到预测结果'}), 404
        
        # 转换为DataFrame
        data = []
        for result in results:
            data.append({
                '业务板块': result.company_name,
                '预测年份': result.prediction_year,
                '销售收入(万元)': result.predicted_sales_amount,
                '净利润(万元)': result.predicted_unit_profit,
                '净利润率(%)': (result.predicted_unit_profit / result.predicted_sales_amount * 100) if result.predicted_sales_amount > 0 else 0,
                '置信度(%)': result.confidence_score * 100,
                '影响因素分析': result.influence_factors_analysis
            })
        
        df = pd.DataFrame(data)
        
        # 创建Excel文件
        output = io.BytesIO()
        with pd.ExcelWriter(output, engine='openpyxl') as writer:
            # 写入主要数据
            df.to_excel(writer, sheet_name='预测结果', index=False)
            
            # 创建汇总表
            summary_data = []
            years = sorted(df['预测年份'].unique())
            sectors = ['人药制剂板块', '人药原料板块', '植保板块', '动保板块', '集团合计']
            
            for sector in sectors:
                row = {'业务板块': sector}
                for year in years:
                    sector_year_data = df[(df['业务板块'] == sector) & (df['预测年份'] == year)]
                    if not sector_year_data.empty:
                        row[f'{year}年收入(万元)'] = sector_year_data['销售收入(万元)'].iloc[0]
                        row[f'{year}年利润(万元)'] = sector_year_data['净利润(万元)'].iloc[0]
                    else:
                        row[f'{year}年收入(万元)'] = 0
                        row[f'{year}年利润(万元)'] = 0
                summary_data.append(row)
            
            summary_df = pd.DataFrame(summary_data)
            summary_df.to_excel(writer, sheet_name='汇总表', index=False)
        
        output.seek(0)
        
        # 生成文件名
        filename = f"预测结果_{task.task_name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.xlsx"
        
        return send_file(
            output,
            as_attachment=True,
            download_name=filename,
            mimetype='application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'
        )
        
    except Exception as e:
        return jsonify({'code': 500, 'message': f'导出失败: {str(e)}'}), 500
